import matplotlib.pyplot as plt
from sklearn.datasets import load_boston
from sklearn.linear_model import Lasso
import numpy as np

fig, axs = plt.subplots(2,2, figsize=(16,7))

def plot(data2D, target1D, predict1D, i, labels):
    axs[int(i/2),i%2].plot(data2D, target1D, 'y.', markersize=6,
                           label='Samples')
    axs[int(i/2),i%2].plot(data2D, predict1D, 'k-', label='Prediction')
    axs[int(i/2),i%2].legend()
    axs[int(i/2),i%2].set_title('Boston Dataset, '+ labels[i])
    axs[int(i/2),i%2].grid()
    return

boston = load_boston()
y = boston.target      # Regression fit methods need a target vector
CRIM = 0; ZN = 1; ROOMS = 5; LSTAT = 12
labels = ('Per capita crime','Residential land zoned','Number of
             rooms','Lower status of population')

for i, attribute in zip(range(4),[CRIM, ZN, ROOMS, LSTAT]):
    X = boston.data[:,attribute]   # Representative attributes
    # Regression fit methods need a data matrix as argument
    X = X[:, np.newaxis]           
    lasso = Lasso(alpha=0.3)
    lasso.fit(X, y)
    print("Slope: "+str(lasso.coef_)+", intercept:
          "+str(lasso.intercept_))
    y_pred = lasso.predict(X)
    plot(X, y, y_pred, i, labels)
